package main

import (
	"bufio"
	"bytes"
	"crypto/des"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/go-ping/ping"
	"github.com/pion/stun"
	"gitlab.com/NebulousLabs/go-upnp"
)

const desKey = "ourheyday-yuandaima"

var defaultStunServers = []string{
	"222.137.252.18", "103.69.128.47", "49.0.239.69", "223.244.91.19",
	"103.219.34.70", "103.219.34.72", "8.218.148.237", "125.70.231.96",
	"218.22.52.100", "39.153.135.73", "59.44.218.162", "119.117.142.4",
	"60.14.52.186", "221.7.195.52", "111.53.11.104",
}

func isPrivateIP(ip net.IP) bool {
	ip4 := ip.To4()
	if ip4 == nil {
		return false
	}
	switch {
	case ip4[0] == 10:
		return true
	case ip4[0] == 172 && ip4[1]&0xf0 == 16:
		return true
	case ip4[0] == 192 && ip4[1] == 168:
		return true
	}
	return false
}

func getLocalIPAndMAC() (ipStr, macStr string, err error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return
	}
	for _, iface := range ifaces {
		if (iface.Flags&net.FlagUp == 0) ||
			(iface.Flags&net.FlagLoopback != 0) ||
			(len(iface.HardwareAddr) == 0) {
			continue
		}
		addrs, _ := iface.Addrs()
		for _, a := range addrs {
			ipNet, ok := a.(*net.IPNet)
			if !ok {
				continue
			}
			ip := ipNet.IP.To4()
			if ip == nil {
				continue
			}
			return ip.String(), iface.HardwareAddr.String(), nil
		}
	}
	return "", "", fmt.Errorf("no usable ip or mac")
}

func getSN() string {
	content, err := ioutil.ReadFile("/proc/cmdline")
	if err == nil {
		fields := strings.Fields(string(content))
		for _, field := range fields {
			if strings.HasPrefix(field, "snum=") {
				sn := strings.TrimPrefix(field, "snum=")
				if sn != "" {
					return sn
				}
			}
		}
	}
	
	machineId, err := ioutil.ReadFile("/etc/machine-id")
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(machineId))
}

func tryUPNP() (string, error) {
	client, err := upnp.Discover()
	if err != nil {
		return "", err
	}
	ip, err := client.ExternalIP()
	if err != nil {
		return "", err
	}
	return ip, nil
}

func udpAddr(ip string) *net.UDPAddr {
	addr := net.JoinHostPort(ip, "3478")
	u, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return &net.UDPAddr{IP: net.ParseIP("222.137.252.18"), Port: 3478}
	}
	return u
}

func getMappedAddr(stunServer string, localIP net.IP) (net.IP, int, error) {
	conn, err := net.DialUDP("udp", nil, udpAddr(stunServer))
	if err != nil {
		return nil, 0, err
	}
	defer conn.Close()
	req := stun.MustBuild(stun.TransactionID, stun.BindingRequest)
	_, err = conn.Write(req.Raw)
	if err != nil {
		return nil, 0, err
	}
	conn.SetReadDeadline(time.Now().Add(3 * time.Second))
	buf := make([]byte, 1500)
	n, _, err := conn.ReadFrom(buf)
	if err != nil {
		return nil, 0, err
	}
	var res stun.Message
	res.Raw = buf[:n]
	if err := res.Decode(); err != nil {
		return nil, 0, err
	}
	var xor stun.XORMappedAddress
	if err := xor.GetFrom(&res); err != nil {
		return nil, 0, err
	}
	return xor.IP, xor.Port, nil
}

func detectNATType(localIP string, stunServers []string) (string, string, error) {
	if len(stunServers) < 2 {
		return "", "", fmt.Errorf("need at least 2 stun servers")
	}
	rand.Shuffle(len(stunServers), func(i, j int) { stunServers[i], stunServers[j] = stunServers[j], stunServers[i] })
	lip := net.ParseIP(localIP)
	// direct binding
	mappedIP, mappedPort, err := getMappedAddr(stunServers[0], lip)
	if err != nil {
		return "", "", fmt.Errorf("stun failed: %v", err)
	}
	if mappedIP.Equal(lip) {
		return "NAT0", mappedIP.String(), nil
	}
	// second stun
	mappedIP2, mappedPort2, err := getMappedAddr(stunServers[1], lip)
	if err != nil {
		return "", "", fmt.Errorf("second stun failed")
	}
	if mappedIP.Equal(mappedIP2) && mappedPort == mappedPort2 {
		return "NAT1", mappedIP.String(), nil
	}
	if mappedIP.Equal(mappedIP2) && mappedPort != mappedPort2 {
		return "NAT3", mappedIP.String(), nil
	}
	if !mappedIP.Equal(mappedIP2) {
		return "NAT4", mappedIP.String(), nil
	}
	return "NAT2", mappedIP.String(), nil
}

func detectIPv6() string {
	ipv6Target := "240c::6666"
	pinger, err := ping.NewPinger(ipv6Target)
	if err != nil {
		return "0"
	}
	pinger.Count = 3
	pinger.SetPrivileged(true)
	pinger.SetNetwork("udp6")
	if err = pinger.Run(); err != nil {
		return "0"
	}
	stats := pinger.Statistics()
	if stats.PacketLoss < 100 {
		return "1"
	}
	return "0"
}

func pkcs5Padding(src []byte, blockSize int) []byte {
	padLen := blockSize - len(src)%blockSize
	padding := bytes.Repeat([]byte{byte(padLen)}, padLen)
	return append(src, padding...)
}

func pkcs5Unpadding(src []byte) []byte {
	l := len(src)
	if l == 0 {
		return src
	}
	pad := int(src[l-1])
	return src[:l-pad]
}

func desEncrypt(plain string) (string, error) {
	key := []byte(desKey)[:8]
	block, err := des.NewCipher(key)
	if err != nil {
		return "", err
	}
	src := pkcs5Padding([]byte(plain), block.BlockSize())
	dst := make([]byte, len(src))
	for bs, be := 0, block.BlockSize(); bs < len(src); bs, be = bs+block.BlockSize(), be+block.BlockSize() {
		block.Encrypt(dst[bs:be], src[bs:be])
	}
	return base64.StdEncoding.EncodeToString(dst), nil
}

func desDecrypt(cipherText string) (string, error) {
	key := []byte(desKey)[:8]
	block, err := des.NewCipher(key)
	if err != nil {
		return "", err
	}
	decoded, err := base64.StdEncoding.DecodeString(cipherText)
	if err != nil {
		return "", err
	}
	decrypted := make([]byte, len(decoded))
	for bs, be := 0, block.BlockSize(); bs < len(decoded); bs, be = bs+block.BlockSize(), be+block.BlockSize() {
		block.Decrypt(decrypted[bs:be], decoded[bs:be])
	}
	unpadded := pkcs5Unpadding(decrypted)
	return string(unpadded), nil
}

func loadStunServersFromURL(url string) ([]string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return parseStunList(resp.Body), nil
}

func parseStunList(r io.Reader) []string {
	srvs := []string{}
	buf := bufio.NewScanner(r)
	for buf.Scan() {
		line := strings.TrimSpace(buf.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		ip := line
		if ix := strings.Index(ip, ":"); ix > 0 {
			ip = ip[:ix]
		}
		srvs = append(srvs, ip)
	}
	return srvs
}

func checkBindIdFlag() bool {
	flagPath := "/etc/bindid.flag"
	content, err := ioutil.ReadFile(flagPath)
	if err != nil {
		fmt.Printf("[WARN] Flag file %s not found or cannot be read: %v\n", flagPath, err)
		return false
	}
	
	stripped := strings.TrimSpace(string(content))
	if len(stripped) < 5 {
		fmt.Printf("[WARN] Flag file %s content too short (need >=5 chars, got %d): '%s'\n", flagPath, len(stripped), stripped)
		return false
	}
	
	fmt.Printf("[INFO] Flag file %s valid, content length: %d\n", flagPath, len(stripped))
	return true
}

func main() {
	apiURL := flag.String("api-url", "", "API endpoint url (required)")
	interval := flag.Int("interval", 5, "heartbeat interval (minutes, default 5)")
	stunUrl := flag.String("stun-url", "", "optional remote stun list url (plain text, 1 IP per line)")
	debug := flag.Bool("debug", false, "enable debug output (default false)")
	flag.Parse()

	if *apiURL == "" {
		fmt.Println("[FATAL] --api-url required")
		os.Exit(2)
	}
	sn := getSN()
	if sn == "" {
		fmt.Println("[FATAL] sn acquire failed from /proc/cmdline or /etc/machine-id")
		os.Exit(2)
	}

	stunServers := defaultStunServers
	if *stunUrl != "" {
		if list, err := loadStunServersFromURL(*stunUrl); err == nil && len(list) >= 2 {
			stunServers = list
			fmt.Printf("[INFO] stun list loaded: %v\n", stunServers)
		} else {
			fmt.Println("[WARN] stun-url fetch failed or insufficient, using default")
		}
	}

	for {
		if !checkBindIdFlag() {
			fmt.Printf("[INFO] Skipping this cycle due to invalid flag file, sleeping %d min\n", *interval)
			time.Sleep(time.Duration(*interval) * time.Minute)
			continue
		}
		
		inIP, mac, err := getLocalIPAndMAC()
		if err != nil {
			inIP, mac = "", ""
		}
		outIP, natType := "", ""
		upnpStart := "0"
		if wip, err := tryUPNP(); err == nil && wip != "" {
			natType = "NAT0"
			upnpStart = "1"
			outIP = wip
			fmt.Println("[INFO] UPNP Detected, NAT0, outIP:", outIP)
		} else {
			natType, outIP, err = detectNATType(inIP, stunServers)
			if err != nil {
				natType = ""
				outIP = ""
				fmt.Println("[WARN] NAT type detection failed")
			}
		}
		ipv6Start := detectIPv6()
		data := map[string]interface{}{
			"natType":   natType,
			"outIp":     outIP,
			"inIp":      inIP,
			"mac":       mac,
			"ohdCode":   sn,
			"sn":        sn,
			"upnpStart": upnpStart,
			"ipv6start": ipv6Start,
		}
		bodyJson, _ := json.Marshal(data)
		encData, err := desEncrypt(string(bodyJson))
		if err != nil {
			fmt.Println("[FATAL] data encrypt fail", err)
			os.Exit(1)
		}
		reqPack := map[string]interface{}{
			"result": true,
			"data":   encData,
		}
		bodyReq, _ := json.Marshal(reqPack)

		if *debug {
			fmt.Println("[DEBUG] Data plaintext:", string(bodyJson))
			fmt.Println("[DEBUG] Data encrypted (Base64):", encData)
		}

		req, err := http.NewRequest("POST", *apiURL, bytes.NewReader(bodyReq))
		req.Header.Set("Content-Type", "application/json")
		client := &http.Client{Timeout: 10 * time.Second}
		resp, err := client.Do(req)
		if err != nil {
			fmt.Printf("[ERROR] API error: %v, body=%s\n", err, string(bodyReq))
		} else {
			respBytes, _ := ioutil.ReadAll(resp.Body)
			fmt.Printf("[INFO] Push done code=%d, resp= %s\n", resp.StatusCode, string(respBytes))
			if *debug {
				var jr struct {
					Data string `json:"data"`
				}
				if err := json.Unmarshal(respBytes, &jr); err == nil && jr.Data != "" {
					fmt.Println("[DEBUG] Response data (encrypted):", jr.Data)
					if dec, err := desDecrypt(jr.Data); err == nil {
						fmt.Println("[DEBUG] Response data (decrypted):", dec)
					} else {
						fmt.Println("[DEBUG] Response data decrypt error:", err)
					}
				}
			}
			resp.Body.Close()
		}
		fmt.Printf("[INFO] %s sleep %d min\n", time.Now().Format("2006-01-02 15:04:05"), *interval)
		time.Sleep(time.Duration(*interval) * time.Minute)
	}
}


